#include <uapi/linux/ptrace.h>
#include <linux/mm.h>
#include <linux/sched.h>

struct alloc_info_t {
        u64 size;
        u64 timestamp_ns;
        int stack_id;
        u32 pid;
};

struct combined_alloc_info_t {
        u64 total_size;
        u64 number_of_allocs;
        u32 pid;
        char comm[TASK_COMM_LEN];
};

BPF_HASH(sizes, u64);
BPF_HASH(allocs, u64, struct alloc_info_t, 1000000);
BPF_HASH(memptrs, u64, u64);
BPF_STACK_TRACE(stack_traces, 10240);
BPF_HASH(combined_allocs, u64, struct combined_alloc_info_t, 10240);
BPF_HASH(stat, u32);

static inline void update_statistics_add(u64 stack_id, u64 sz) {
        struct combined_alloc_info_t *existing_cinfo;
        struct combined_alloc_info_t cinfo = {0};
        u64* allBytesCount = 0;
        u64 allBytes = 0;

        existing_cinfo = combined_allocs.lookup(&stack_id);
        if (existing_cinfo != 0)
                cinfo = *existing_cinfo;

        u32 pid = bpf_get_current_pid_tgid();

        cinfo.total_size += sz;
        cinfo.number_of_allocs += 1;
        cinfo.pid = pid;
        bpf_get_current_comm(&cinfo.comm, sizeof(cinfo.comm));

        combined_allocs.update(&stack_id, &cinfo);

        allBytesCount = stat.lookup(&pid);
        if (allBytesCount != 0)
                allBytes = *allBytesCount;
        allBytes += sz;
        stat.update(&pid, &allBytes);
}

static inline void update_statistics_del(u64 stack_id, u64 sz) {
        struct combined_alloc_info_t *existing_cinfo;
        struct combined_alloc_info_t cinfo = {0};

        u64* allBytesCount = 0;
        u64 allBytes = 0;

        existing_cinfo = combined_allocs.lookup(&stack_id);
        if (existing_cinfo != 0)
                cinfo = *existing_cinfo;

        if (sz >= cinfo.total_size)
                cinfo.total_size = 0;
        else
                cinfo.total_size -= sz;

        if (cinfo.number_of_allocs > 0)
                cinfo.number_of_allocs -= 1;

        combined_allocs.update(&stack_id, &cinfo);

        u32 pid = bpf_get_current_pid_tgid();
        allBytesCount = stat.lookup(&pid);
        if (allBytesCount != 0)
                allBytes = *allBytesCount;
        allBytes -= sz;
        if (allBytes < 0)
                allBytes = 0;
        stat.update(&pid, &allBytes);
}

static bool isCared()
{
        u64 *ppid_cared = NULL, key = 0 ;
        ppid_cared = stat.lookup(&key);
        u64 pid = bpf_get_current_pid_tgid();
        if (ppid_cared && *ppid_cared == pid)
        {
                return true;
        }

        char comm[TASK_COMM_LEN] = {0};
        bpf_get_current_comm(&comm, sizeof(comm));
        if (strcmp(comm, "COMM_CARED") != 0)
                return false;

        stat.update(&key, &pid);
        return true;
}

inline int kprobe__do_mmap(struct pt_regs *ctx, 
        struct file *file, 
        unsigned long addr,
        u64 len
        ) {

        if (!isCared()) return 0;
        u64 pid = bpf_get_current_pid_tgid();

        if (SHOULD_PRINT)
                bpf_trace_printk("alloc entered start, size = %u, pid %ld\\n", 
                        len, (u32)pid);

        u64 size64 = len;
        sizes.update(&pid, &size64);

        if (SHOULD_PRINT)
                bpf_trace_printk("alloc entered, size = %u\\n", len);
        return 0;
}

static inline int gen_alloc_exit2(struct pt_regs *ctx, u64 address) {
        u64 pid = bpf_get_current_pid_tgid();

        u64* size64 = sizes.lookup(&pid);
        struct alloc_info_t info = {0};

        if (size64 == 0)
                return 0; // missed alloc entry

        info.size = *size64;
        info.pid = (u32)pid;
        sizes.delete(&pid);

        if (address != 0) {
                info.timestamp_ns = bpf_ktime_get_ns();
                info.stack_id = stack_traces.get_stackid(ctx, STACK_FLAGS);
                allocs.update(&address, &info);
                update_statistics_add(info.stack_id, info.size);
        }

        if (SHOULD_PRINT) {
                bpf_trace_printk("alloc exited, size = %lu, result = %lx\\n",
                                 info.size, address);
        }
        return 0;
}

inline int kretprobe__do_mmap(struct pt_regs *ctx) {
        if (!isCared()) return 0;
        return gen_alloc_exit2(ctx, PT_REGS_RC(ctx));
}

inline int kprobe__do_munmap(struct pt_regs *ctx, 
        struct mm_struct *mm,
        u64 address,
        size_t len
        ) {

        if (!isCared()) return 0;

        u64 pid = bpf_get_current_pid_tgid();

        u64 addr = (u64)address;
        struct alloc_info_t *info = allocs.lookup(&addr);
        if (info == 0)
                return 0;

        allocs.delete(&addr);
        update_statistics_del(info->stack_id, info->size);

        if (SHOULD_PRINT) {
                bpf_trace_printk("free entered, address = %lx, size = %lu\\n",
                                 address, info->size);
        }
        return 0;
}